import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from sklearn.tree import DecisionTreeRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score

# 读取包含法律案件的数据集
df = pd.read_csv('D:/python/Lab1/Last.csv')

# 将日期转换为年份和月份两个新特征
df['year'] = pd.to_datetime(df['start_time']).dt.year
df['month'] = pd.to_datetime(df['start_time']).dt.month

# 对 案件原因 进行独热编码
dummy_df_1 = pd.get_dummies(df['案件基本情况'], prefix='base')
dummy_df_2 = pd.get_dummies(df['案件原因'], prefix='cause')

# 将独热编码后的数据与其他特征合并
features_df = pd.concat([df[['year', 'month']], dummy_df_1, dummy_df_2], axis=1)

# 统计每个案件原因在数据集中出现的总次数
cause_counts = df['案件原因'].value_counts()

# 计算每个案件原因在所有案件中所占的比例
total_cases = len(df)
case_ratios = cause_counts / total_cases

# 将 case_ratio 添加到 DataFrame 中
df = df.merge(case_ratios.rename('case_ratio'), left_on='案件原因', right_index=True)

# 输出结果
print(df)

# 划分数据集为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(features_df, df['case_ratio'],
                                                    test_size=0.2, random_state=42)

# 创建决策树回归模型，并训练模型
model = DecisionTreeRegressor(random_state=42)
model.fit(X_train, y_train)

# 对测试集进行预测
y_pred = model.predict(X_test)
# 计算MAE、RMSE和MSE
mae = mean_absolute_error(y_test, y_pred)
rmse = np.sqrt(mean_squared_error(y_test, y_pred))
mse = mean_squared_error(y_test, y_pred)

print("MAE:", mae)
print("RMSE:", rmse)
print("MSE:", mse)

# 创建绘图对象和子图
fig, ax = plt.subplots(figsize=(8, 6))

# 绘制柱状图
ax.bar(["MAE", "RMSE", "MSE"], [mae, rmse, mse])

# 添加标签和标题
ax.set_title("Model Performance on Test Set")
ax.set_ylabel("Performance Metric Value")
ax.set_xlabel("Performance Metric")

plt.show()

